import os
from common.benchmark_var import GlobalEnv
from utils.executor import Executor
from common.benchmark_var import MeasureUnit
from workload.workloads_manager import WorkloadManager
from framework.native_workload import NativeWorkload
from framework.java_workload import JavaWorkload
from copy import deepcopy
import shutil

class CUDAWorkload(JavaWorkload):
    def __init__(self, name, options, baseline):
        super().__init__(name, options, baseline)
        self.name = ''
        self.lang = 'C'
        self.description = ''
        self.result_file_name = ''
        self.desc = ''
        self.program_options = ""
        self.exe = ""
        self.script = "cuda.sh"
        self.valid = True
        self.support_multi_instance = False
        
    def setup(self):
        self._copy_src2build()
        executor = Executor(desc='make_cuda',
                        command='make -j',
                        exec_dir=self.build_dir,
                        debug=self.debug)
        res = executor.exec()

    def _copy_src2build(self):
        files = os.listdir(self.src_dir)
        for file in files:
            src_file = os.path.join(self.src_dir, file)
            dst_file = os.path.join(self.build_dir, file)
            if os.path.exists(dst_file):
                shutil.rmtree(dst_file) if os.path.isdir(dst_file) else os.remove(dst_file)
            if os.path.isdir(src_file):
                shutil.copytree(src_file, dst_file)
            else:
                shutil.copy(src_file, self.build_dir)

    def get_instance_cmd(self, instance):
        program = os.path.join(self.build_dir, self.script)
        executors = []
        cmd = "bash {} {}".format(program, self.program_options)

        executor = Executor(desc=self.desc,
                            command=cmd,
                            debug=self.debug,
                            exec_dir=instance.instance_dir)
        executors.append(executor)
        return executors
    
    def get_raw_score(self, instance):
        result_file = os.path.join(instance.instance_dir, self.result_file_name)
        if not os.path.exists(result_file):
            return 0
        with open(result_file, 'r') as f:
            lines = f.readlines()
            for line in lines:
                line = line.strip()
                if "COUNT" in line:
                    str_list = line.split('|')
                    score = float(str_list[1])
                    return score
        return 0

    def setup_dir(self):
        self.work_dir = os.path.join(GlobalEnv.GCBS_WORKLOAD_DIR, "cuda")
        self.src_dir = os.path.join(self.work_dir, 'src')
        self.build_dir = os.path.join(self.work_dir, 'build')
        self.data_dir = os.path.join(self.work_dir, 'data')
        self.log_dir = os.path.join(self.work_dir, 'log')
        self.bin_dir = os.path.join(self.work_dir, 'bin')
        self.run_dir = os.path.join(self.work_dir, 'run', self.name)
        dir_list = [self.work_dir, self.src_dir, self.build_dir, self.data_dir,
                    self.log_dir, self.run_dir, self.bin_dir]
        for directory in dir_list:
            if not os.path.exists(directory):
                os.makedirs(directory)


@WorkloadManager.register_workload('bs-download-cuda')
class CUDABsDownloadSpeedWorkload(CUDAWorkload):
    DESC = "CUDA: bs-download-speed"
    def __init__(self, name, options, baseline):
        super().__init__(name, options, baseline)
        self.name = "bs-download-cuda"
        self.description = CUDABsDownloadSpeedWorkload.DESC
        self.unit_of_measure = MeasureUnit.SCORE
        self.result_file_name = "bs-download-speed.log"
        self.program_options = "-b bs-download-speed"
        self.setup_dir()
        self.desc = "bs-download-speed"

@WorkloadManager.register_workload('bs-readback-cuda')
class CUDABsReadbackSpeedWorkload(CUDAWorkload):
    DESC = "CUDA: bs-readback-speed"
    def __init__(self, name, options, baseline):
        super().__init__(name, options, baseline)
        self.name = "bs-readback-cuda"
        self.description = CUDABsReadbackSpeedWorkload.DESC
        self.unit_of_measure = MeasureUnit.SCORE
        self.result_file_name = "bs-readback-speed.log"
        self.program_options = "-b bs-readback-speed"
        self.setup_dir()
        self.desc = "bs-readback-speed"

@WorkloadManager.register_workload('gmem-rw-cuda')
class CUDAGmemRwBandwidthWorkload(CUDAWorkload):
    DESC = "CUDA: gmem-rw-bandwidth"
    def __init__(self, name, options, baseline):
        super().__init__(name, options, baseline)
        self.name = "gmem-rw-cuda"
        self.description = CUDAGmemRwBandwidthWorkload.DESC
        self.unit_of_measure = MeasureUnit.SCORE
        self.result_file_name = "gmem-rw-bandwidth.log"
        self.program_options = "-b gmem-rw-bandwidth"
        self.setup_dir()
        self.desc = "gmem-rw-bandwidth"

@WorkloadManager.register_workload('maxhalf-cuda')
class CUDAMaxhalfComputeWorkload(CUDAWorkload):
    DESC = "CUDA: maxhalf-compute"
    def __init__(self, name, options, baseline):
        super().__init__(name, options, baseline)
        self.name = "maxhalf-cuda"
        self.description = CUDAMaxhalfComputeWorkload.DESC
        self.unit_of_measure = MeasureUnit.SCORE
        self.result_file_name = "maxhalf-compute.log"
        self.program_options = "-b maxhalf-compute"
        self.setup_dir()
        self.desc = "maxhalf-compute"

@WorkloadManager.register_workload('maxint-cuda')
class CUDAMaxintComputeWorkload(CUDAWorkload):
    DESC = "CUDA: maxint-compute"
    def __init__(self, name, options, baseline):
        super().__init__(name, options, baseline)
        self.name = "maxint-cuda"
        self.description = CUDAMaxintComputeWorkload.DESC
        self.unit_of_measure = MeasureUnit.SCORE
        self.result_file_name = "maxint-compute.log"
        self.program_options = "-b maxint-compute"
        self.setup_dir()
        self.desc = "maxint-compute"

@WorkloadManager.register_workload('maxsp-cuda')
class CUDAMaxspComputeWorkload(CUDAWorkload):
    DESC = "CUDA: maxsp-compute"
    def __init__(self, name, options, baseline):
        super().__init__(name, options, baseline)
        self.name = "maxsp-cuda"
        self.description = CUDAMaxspComputeWorkload.DESC
        self.unit_of_measure = MeasureUnit.SCORE
        self.result_file_name = "maxsp-compute.log"
        self.program_options = "-b maxsp-compute"
        self.setup_dir()
        self.desc = "maxsp-compute"

@WorkloadManager.register_workload('maxdp-cuda')
class CUDAMaxdpComputeWorkload(CUDAWorkload):
    DESC = "CUDA: maxdp-compute"
    def __init__(self, name, options, baseline):
        super().__init__(name, options, baseline)
        self.name = "maxdp-cuda"
        self.description = CUDAMaxdpComputeWorkload.DESC
        self.unit_of_measure = MeasureUnit.SCORE
        self.result_file_name = "maxdp-compute.log"
        self.program_options = "-b maxdp-compute"
        self.setup_dir()
        self.desc = "maxdp-compute"

@WorkloadManager.register_workload('fft-test-cuda')
class CUDAFFTTestWorkload(CUDAWorkload):
    DESC = "CUDA: fft-test"
    def __init__(self, name, options, baseline):
        super().__init__(name, options, baseline)
        self.name = "fft-test-cuda"
        self.description = CUDAFFTTestWorkload.DESC
        self.unit_of_measure = MeasureUnit.SCORE
        self.result_file_name = "fft-test.log"
        self.program_options = "-b fft-test"
        self.setup_dir()
        self.desc = "fft-test"

@WorkloadManager.register_workload('gemm-test-cuda')
class CUDAGEMMTestWorkload(CUDAWorkload):
    DESC = "CUDA: gemm-test"
    def __init__(self, name, options, baseline):
        super().__init__(name, options, baseline)
        self.name = "gemm-test-cuda"
        self.description = CUDAGEMMTestWorkload.DESC
        self.unit_of_measure = MeasureUnit.SCORE
        self.result_file_name = "gemm-test.log"
        self.program_options = "-b gemm-test"
        self.setup_dir()
        self.desc = "gemm-test"




